""" Full assembly of the parts to form the complete network """

from .resunet_parts import *


class ResUNet(nn.Module):
    def __init__(self, n_channels, n_classes, device,bilinear=False,classification =True):
        super(ResUNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear
        self.classification = classification
        self.device = device
        factor = 2 if bilinear else 1
        
        self.inc1 = DoubleConv_res(n_channels, 64//2)
        self.down1_1 = Down_res(64//2, 128//2)
        self.down1_2 = Down_res(128//2, 256//2)
        self.down1_3 = Down_res(256//2, 512//2)
        self.down1_4 = Down_res(512//2, 1024//2 // factor)

        self.inc2 = DoubleConv_res(n_channels, 64//2)
        self.down2_1 = Down_res(64//2, 128//2)
        self.down2_2 = Down_res(128//2, 256//2)
        self.down2_3 = Down_res(256//2, 512//2)
        self.down2_4 = Down_res(512//2, 1024//2 // factor)


        self.up1 = Up_res(1024, 512 // factor, bilinear)
        self.up2 = Up_res(512, 256 // factor, bilinear)
        self.up3 = Up_res(256, 128 // factor, bilinear)
        self.up4 = Up_res(128, 64, bilinear)
        self.outc = OutConv(64, n_classes)



    def forward(self, image1, image2):   
        # encoder for image1
        x1_1 = self.inc1(image1)
        x1_2 = self.down1_1(x1_1)
        x1_3 = self.down1_2(x1_2)
        x1_4 = self.down1_3(x1_3)
        x1_5 = self.down1_4(x1_4)

        # encoder for image2
        x2_1 = self.inc2(image2)
        x2_2 = self.down2_1(x2_1)
        x2_3 = self.down2_2(x2_2)
        x2_4 = self.down2_3(x2_3)
        x2_5 = self.down2_4(x2_4)

        # feature concat for images
        x1 = torch.cat([x1_1,x2_1], dim=1)
        x2 = torch.cat([x1_2,x2_2], dim=1)
        x3 = torch.cat([x1_3,x2_3], dim=1)
        x4 = torch.cat([x1_4,x2_4], dim=1)
        x5 = torch.cat([x1_5,x2_5], dim=1)

        # decoder
        x_up1 = self.up1(x5, x4)
        x_up2 = self.up2(x_up1, x3)
        x_up3 = self.up3(x_up2, x2)
        x = self.up4(x_up3, x1)
        logits = self.outc(x)

        logits =  torch.sigmoid(logits)
        return logits

